from joblib import load
import numpy as np

class Classification(object):
    def __init__(self):
        self.knn = load('knn_classify.joblib')
        self.scaler = load('feature_scaler.joblib')
        self.le = load('label_encoder.joblib')

    def get_predictor_category(self, new_data):
        """
        获取分类结果
        """
        # 转换为数组并保持特征顺序
        feature_order = [
            'eps', 'total_revenue_ps', 'undist_profit_ps', 'gross_margin',
            'fcff', 'fci', 'tangible_asset', 'bps', 'grossprofit_margin', 'npta'
        ]
        new_values = np.array([new_data[col] for col in feature_order])

        # 标准化新数据
        new_scaled = self.scaler.transform(new_values)

        # 预测分类
        predicted_label = self.knn.predict(new_scaled)
        predicted_category = self.le.inverse_transform(predicted_label)
        return predicted_category[0]

if __name__ == '__main__':
    ci = Classification()
    new_data = {
        "eps": '-0.8309',
        "total_revenue_ps": '11.9673',
        "undist_profit_ps": '7.0209',
        "gross_margin": '11586700000',
        "fcff": '1287910000',
        "fci": '526340000',
        "tangible_asset": '179873000000',
        "bps": '20.2568',
        "grossprofit_margin": '8.1152',
        "npta": '-0.5821',
    }
    predicted_category = ci.get_predictor_category(new_data)
    print(predicted_category)